# Author: miaokl
# Date: 2023-07-06
import traceback
from datetime import datetime, timezone

from flask_restx import Resource, marshal_with
from sqlalchemy import and_
from flask_jwt_extended import (
    create_access_token, create_refresh_token, get_jwt
)

from app.utils.auth import jwt_required
from app.api import register_scheme
from app.scheme import response, user as _user
from app.models import db, UserInfo, TokenBlockList
from app.utils.common import paginate, do_md5
from app.parser import get_parser
from app.utils.validation import validate, blank_not_allowed
from app.utils.exception import BaseError, AddRecordException, ModifyRecordException, DelRecordException
from app.utils.db import DataBaseHelper


ns = register_scheme(__file__)


@ns.expect(get_parser('auth', 'auth_rules'))
@ns.route('/user')
class UserManagement(Resource):
    """
    Add user
    Get users
    """
    decorators = [jwt_required()]

    @marshal_with(_user.users)
    @ns.expect(get_parser('page', 'page_rules'))  # for verify pagination params in api doc 
    def get(self):
        return response({
            'total': UserInfo.queryset().count(),
            'list': paginate(UserInfo.queryset()),
        })

    @marshal_with(_user.user_out)
    @ns.expect(_user.user_in, validate=True)
    @validate([  # User-defined verification scheme
        {
            "field": ['username', 'nickname', 'password'],
            'validator': blank_not_allowed,
        }
    ])
    def post(self):
        payload = ns.payload

        _username = payload.get('username')

        data = {
            'username': _username,
            'nickname': payload.get('nickname'),
            'password': do_md5(payload.get('password')),
        }

        # Check wheter current user exists
        user_check = UserInfo.queryset().filter(UserInfo.username == _username).count()

        if user_check > 0:
            return BaseError(msg='当前用户已存在')

        with db.trans():
            _id = DataBaseHelper.execute_create(
                'user_info',
                data=data,
            )

            if not _id:
                raise AddRecordException()

            return response({
                'id': _id,
                **data
            })


@ns.expect(get_parser('auth', 'auth_rules'))
@ns.route("/user/<int:uid>")
class UserManagementDetail(Resource):
    """
    Delete user
    Modify user
    """
    decorators = [jwt_required()]

    @marshal_with(_user.user_out)
    def delete(self, uid):
        with db.trans():
            # The last user cannot be deleted
            user_count = UserInfo.queryset().count()
            if user_count == 1:
                return BaseError(msg='不允许删除所有用户')
            rows = DataBaseHelper.execute_delete(
                'user_info',
                where={
                    'id': uid,
                },
                logic=True
            )
            if not rows:
                raise DelRecordException()
        return response({})
    
    @marshal_with(_user.user_out)
    @ns.expect(_user.user_in_without_passwd, validate=True)
    @validate([
        {
            "field": ['username', 'nickname'],
            'validator': blank_not_allowed,
        }
    ])
    def put(self, uid):
        payload = ns.payload

        _username = payload.get('username')

        data = {
            'username': _username,
            'nickname': payload.get('nickname'),
        }

        # Determine if the user is repeated except for itself
        user_count = UserInfo.queryset().filter(and_(
            UserInfo.id != uid,
            UserInfo.username == _username,
        )).count()
        if user_count > 0:
            return BaseError(msg='当前用户已存在')

        with db.trans():
            rows = DataBaseHelper.execute_update(
                'user_info',
                data=data,
                where={
                    'id': uid
                }
            )
            if not rows:
                raise ModifyRecordException()
            return response({
                'id': uid,
                **data
            })


@ns.route("/user/login")
class UserLogin(Resource):

    @marshal_with(_user.user_out_login)
    @ns.expect(_user.user_in_login, validate=True)
    @validate([
        {
            "field": ['username', 'password'],
            'validator': blank_not_allowed,
        }
    ])
    def post(self):
        payload = ns.payload

        _username = payload.get('username')
        _password = payload.get('password')

        # Check whether the user exists
        _user = UserInfo.queryset().filter(UserInfo.username == _username).first()

        if not _user:
            return BaseError(msg='当前用户不存在')

        # Check whether the user password is correct
        if do_md5(_password) != _user.password:
            return BaseError(msg='用户密码不正确')

        # Generate access_token and refresh_token
        access_token = create_access_token(identity={
            'user_id': _user.id,
        })
        refresh_token = create_refresh_token(identity={
            'user_id': _user.id,
        })

        return response({
            'token': access_token,
            'refresh_token': refresh_token,
        })


@ns.expect(get_parser('auth', 'auth_rules'))
@ns.route("/user/logout")
class UserLogout(Resource):
    decorators = [jwt_required()]

    def post(self):
        jti = get_jwt()["jti"]
        now = datetime.now(timezone.utc)
        db.session.add(TokenBlockList(jti=jti, created_at=now))
        db.session.commit()

        return response({})
